import socket
import os
import struct
from typing import Callable, Optional

class FileSender:
    PROTOCOL_VERSION = 0x01
    BUFFER_SIZE = 8192  # 8KB
    SOCKET_TIMEOUT = 30  # 30秒
    FILTER_FOLDER_LIST = (".obsidian",)
    FILTER_FILE_LIST = ("",)

    def __init__(self):
        self.progress_callback: Optional[Callable] = None

    def set_progress_callback(self, callback: Callable):
        """设置进度回调函数"""
        self.progress_callback = callback

    def _safe_callback(self, callback_type: str, *args):
        """安全地调用回调函数"""
        if self.progress_callback:
            self.progress_callback(callback_type, *args)

    def _count_files(self, folder_path: str) -> int:
        """计算文件夹中的文件总数"""
        total_files = 0
        for root, _, files in os.walk(folder_path):
            total_files += len(files)
        return total_files

    def _send_folder_header(self, conn: socket.socket, folder_path: str):
        """发送文件夹头部信息"""
        # 发送协议版本
        conn.send(bytes([self.PROTOCOL_VERSION]))
        
        # 发送文件夹名称
        folder_name = os.path.basename(folder_path)
        name_bytes = folder_name.encode('utf-8')
        conn.send(struct.pack('>H', len(name_bytes)))  # 发送字符串长度（2字节）
        conn.send(name_bytes)
        
        # 发送文件总数
        total_files = self._count_files(folder_path)
        conn.send(struct.pack('>Q', total_files))  # 8字节long
        
        self._safe_callback('total_files', total_files)

    def _send_single_file(self, conn: socket.socket, file_path: str, relative_path: str):
        """发送单个文件"""
        # 发送文件标识符
        conn.send(bytes([0x02]))
        
        # 发送相对路径
        path_bytes = relative_path.encode('utf-8')
        conn.send(struct.pack('>H', len(path_bytes)))
        conn.send(path_bytes)
        
        # 发送文件大小
        file_size = os.path.getsize(file_path)
        conn.send(struct.pack('>Q', file_size))
        
        # 发送文件内容
        transferred = 0
        with open(file_path, 'rb') as f:
            while True:
                data = f.read(self.BUFFER_SIZE)
                if not data:
                    break
                conn.send(data)
                transferred += len(data)
                self._safe_callback('file_progress', relative_path, transferred, file_size)

    def _traverse_folder(self, conn: socket.socket, folder_path: str, base_path: str = ""):
        """递归遍历文件夹并发送文件"""
        for item in os.listdir(folder_path):
            full_path = os.path.join(folder_path, item)
            relative_path = os.path.join(base_path, item).replace('\\', '/')
            if os.path.isfile(full_path):
                self._send_single_file(conn, full_path, relative_path)
                # and os.path.basename(full_path) not in self.FILTER_FOLDER_LIST
            elif os.path.isdir(full_path):
                self._traverse_folder(conn, full_path, relative_path)

    def send_folder(self, folder_path: str, port: int = 3000, host: str = '192.168.1.9'):
        """作为客户端连接到服务器并发送文件夹"""
        try:
            print(f"正在连接到服务器：{host}:{port}")
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as client_socket:
                client_socket.connect((host, port))
                client_socket.settimeout(self.SOCKET_TIMEOUT)
                
                print(f"已连接到服务器：{host}:{port}")
                self._send_folder_header(client_socket, folder_path)
                self._traverse_folder(client_socket, folder_path)
                self._safe_callback('complete')
                    
        except Exception as e:
            print(f"连接错误：{str(e)}")
            self._safe_callback('error', str(e))
            raise
'''
if __name__ == '__main__':
    # 使用示例
    def progress_callback(callback_type, *args):
        if callback_type == 'total_files':
            print(f"总文件数：{args[0]}")
        elif callback_type == 'file_progress':
            print(f"文件 {args[0]} 进度：{args[1]}/{args[2]}")
        elif callback_type == 'complete':
            print("传输完成")
        elif callback_type == 'error':
            print(f"错误：{args[0]}")

    sender = FileSender()
    sender.set_progress_callback(progress_callback)
'''
    # 使用示例：sender.send_folder("要发送的文件夹路径", 端口号) 